package com.shell.guard.security;

import java.io.Serializable;
import java.util.HashMap;
import java.util.Map;

import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.store.JwtAccessTokenConverter;

public class CustomTokenEnhancer extends JwtAccessTokenConverter implements Serializable {
	private static final long serialVersionUID = -4802799816735213420L;

	public OAuth2AccessToken enhance(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
		CustomUserDetails userDetails = (CustomUserDetails) authentication.getPrincipal();
		authentication.getUserAuthentication().getPrincipal();
		Map<String, Object> info = new HashMap<String, Object>();
		info.put("X-AOHO-UserId", userDetails.getUserId());
		info.put("X-AOHO-ENT", userDetails.getEnt());
		info.put("TERM_CODE_IN_HEADER", userDetails.getTerminalCode());

		DefaultOAuth2AccessToken customAccessToken = new DefaultOAuth2AccessToken(accessToken);
		customAccessToken.setAdditionalInformation(info);

		OAuth2AccessToken enhancedToken = super.enhance(customAccessToken, authentication);
		enhancedToken.getAdditionalInformation().put("X-AOHO-ClientId", userDetails.getClientId());
		return enhancedToken;
	}
}
